import json
import os
import sys
import time
from typing import Dict, List, Optional

# 确保项目根目录在Python路径中
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入提示词量化评估中的SparkModelEvaluator类
try:
    # 使用绝对导入路径
    from 提示词量化评估.spark_evaluator import SparkModelEvaluator
    print("成功导入SparkModelEvaluator类")
except ImportError as e:
    print(f"警告: 无法导入SparkModelEvaluator类，评估功能将不可用。错误信息: {str(e)}")
    SparkModelEvaluator = None

# 如果直接导入失败，尝试添加具体的提示词量化评估目录到Python路径
try:
    if SparkModelEvaluator is None:
        evaluation_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "提示词量化评估")
        if evaluation_dir not in sys.path:
            sys.path.append(evaluation_dir)
        from spark_evaluator import SparkModelEvaluator
        print("通过添加评估目录到Python路径成功导入SparkModelEvaluator类")
except ImportError as e:
    if SparkModelEvaluator is None:
        print(f"再次尝试导入失败: {str(e)}")

class SparkOptimizer:
    def __init__(self):
        self.evaluator = SparkModelEvaluator() if SparkModelEvaluator else None
        self.optimization_history = []
        self.script_dir = os.path.dirname(os.path.abspath(__file__))
        if self.evaluator:
            print("已成功初始化SparkModelEvaluator，可以使用星火大模型进行提示词优化和评估")
        else:
            print("注意: SparkModelEvaluator初始化失败，大模型优化功能不可用")
    
    def optimize_with_spark(self, long_text: str) -> str:
        """
        使用讯飞星火大模型直接优化提示词
        参数:
            long_text: 用户输入的长文本提示词
        返回:
            优化后的提示词
        """
        if not self.evaluator:
            raise ValueError("错误: SparkModelEvaluator不可用，无法使用星火大模型进行优化")
        
        print("正在使用讯飞星火大模型优化提示词...")
        # 构建用于优化提示词的提示
        optimize_prompt = f"""
        请你作为一名提示词优化专家，帮我优化下面的长文本提示词，使其更简洁、精准、有效。
        优化要求：
        1. 保留核心信息和意图
        2. 提高提示词的清晰度和明确性
        3. 使提示词更符合大模型的理解习惯
        4. 保持简洁，但不要遗漏重要信息
        5. 优化后的提示词应保持原始意图不变
        原始提示词：
        {long_text}
        请直接返回优化后的提示词，不要添加任何额外的解释或说明。
        """
        # 调用星火模型进行优化
        optimized_prompt, _ = self.evaluator.call_spark_model(optimize_prompt, temperature=0.7)
        
        if not optimized_prompt:
            raise ValueError("星火大模型优化失败")
        
        return optimized_prompt.strip()
    
    def optimize_prompt(self, long_text: str) -> Dict[str, str]:
        """
        使用讯飞星火大模型优化长文本提示词
        参数:
            long_text: 用户输入的长文本提示词
        返回:
            包含原始提示词和优化后提示词的字典
        """
        print(f"正在优化提示词: {long_text[:50]}{'...' if len(long_text) > 50 else ''}")
        
        # 使用星火大模型进行优化
        try:
            optimized_prompt = self.optimize_with_spark(long_text)
            optimization_method = "spark_model"
        except Exception as e:
            print(f"优化失败: {str(e)}")
            return None
        
        # 记录优化历史
        optimization_result = {
            "original_prompt": long_text,
            "optimized_prompt": optimized_prompt,
            "optimization_method": optimization_method,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
        }
        self.optimization_history.append(optimization_result)
        return optimization_result
    
    def compare_prompts(self, original_prompt: str, optimized_prompt: str) -> Optional[Dict]:
        """
        比较原始提示词和优化后提示词的效果
        参数:
            original_prompt: 原始长文本提示词
            optimized_prompt: 优化后的提示词
        返回:
            包含比较结果的字典，如果评估器不可用则返回None
        """
        if not self.evaluator:
            print("评估器不可用，无法比较提示词效果")
            return None
        
        print("正在比较原始提示词和优化后提示词的效果...")
        
        # 评估原始提示词
        print("评估原始提示词...")
        original_result = self.evaluator.evaluate_prompt(original_prompt)
        
        # 评估优化后的提示词
        print("评估优化后的提示词...")
        optimized_result = self.evaluator.evaluate_prompt(optimized_prompt)
        
        # 构建比较结果
        comparison = {
            "original_prompt": original_prompt,
            "optimized_prompt": optimized_prompt,
            "original_scores": original_result["avg_scores"] if original_result else None,
            "optimized_scores": optimized_result["avg_scores"] if optimized_result else None,
            "improvement": None
        }
        
        # 计算改进度
        if original_result and optimized_result:
            original_total = original_result["avg_scores"]["total"]
            optimized_total = optimized_result["avg_scores"]["total"]
            comparison["improvement"] = round(optimized_total - original_total, 2)
            
        return comparison
    
    def save_optimization_result(self, result: Dict, filename: str = None):
        """
        保存优化结果到文件
        
        参数:
            result: 优化结果字典
            filename: 保存的文件名，默认为None（自动生成）
        """
        # 如果没有提供文件名，根据提示词和时间戳生成唯一文件名
        if filename is None:
            # 获取当前时间戳，确保唯一性
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            
            # 从结果中提取提示词信息
            if result and "original_prompt" in result:
                prompt_preview = result["original_prompt"][:10] if result["original_prompt"] else "unknown"
            else:
                prompt_preview = "unknown"
            
            # 替换文件名中的非法字符
            prompt_preview = ";".join(c for c in prompt_preview if c.isalnum() or c in ("_", "-", " ")).strip()
            
            # 构建文件名
            filename = f"spark_optimization_result_{prompt_preview}_{timestamp}.json"
        
        # 如果提供的文件名不是绝对路径，则将其放在脚本所在目录下
        if not os.path.isabs(filename):
            filename = os.path.join(self.script_dir, filename)
        
        # 保存结果
        with open(filename, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        
        print(f"星火大模型优化结果已保存到 {filename}")
    
    def print_optimization_result(self, result: Dict):
        """
        打印优化结果
        
        参数:
            result: 优化结果字典
        """
        print("\n" + "="*60)
        print("原始提示词:")
        print(result["original_prompt"][:200] + ("..." if len(result["original_prompt"]) > 200 else ""))
        
        print("\n优化后的提示词:")
        print(result["optimized_prompt"])
        print("="*60 + "\n")
    
    def print_comparison_result(self, comparison: Dict):
        """
        打印提示词比较结果
        
        参数:
            comparison: 比较结果字典
        """
        if not comparison:
            return
        
        print("\n" + "="*60)
        print("提示词效果比较结果")
        print("="*60)
        
        print(f"原始提示词总分: {comparison['original_scores']['total'] if comparison['original_scores'] else 'N/A'}")
        print(f"优化后提示词总分: {comparison['optimized_scores']['total'] if comparison['optimized_scores'] else 'N/A'}")
        
        if comparison["improvement"] is not None:
            improvement_str = f"+{comparison['improvement']}" if comparison["improvement"] > 0 else str(comparison["improvement"])
            print(f"改进度: {improvement_str}")
            
            if comparison["improvement"] > 0:
                print("✓ 优化后的提示词效果更好！")
            elif comparison["improvement"] < 0:
                print("⚠️ 优化后的提示词效果有所下降。")
            else:
                print("优化后的提示词效果与原始提示词相当。")
        
        print("="*60 + "\n")

def main():
    """主函数"""
    optimizer = SparkOptimizer()
    
    print("欢迎使用讯飞星火大模型提示词优化工具")
    print("此工具可以将长文本提示词通过讯飞星火大模型转换为更精准、更有效的提示词。")
    
    if not optimizer.evaluator:
        print("\n警告: 无法初始化SparkModelEvaluator，大模型优化功能不可用。")
        print("请确保已正确配置星火模型连接。")
        return
    
    while True:
        print("\n请选择操作:")
        print("1. 使用讯飞星火大模型优化单个提示词")
        print("2. 比较优化前后的提示词效果")
        print("3. 退出")
        
        choice = input("请输入选择 (1/2/3): ")
        
        if choice == "1":
            long_text = input("请输入要优化的长文本提示词: ")
            result = optimizer.optimize_prompt(long_text)
            if result:
                optimizer.print_optimization_result(result)
                
                # 询问是否保存结果
                save_choice = input("是否保存优化结果？(y/n): ").lower()
                if save_choice == "y":
                    optimizer.save_optimization_result(result)
            else:
                print("优化失败，请重试。")
        
        elif choice == "2":
            if not optimizer.evaluator:
                print("错误: 无法导入评估模块，无法比较提示词效果。")
                continue
            
            long_text = input("请输入要优化并比较的长文本提示词: ")
            
            # 优化提示词
            result = optimizer.optimize_prompt(long_text)
            if not result:
                print("优化失败，无法进行比较。")
                continue
            
            # 打印优化结果
            optimizer.print_optimization_result(result)
            
            # 比较优化结果
            comparison = optimizer.compare_prompts(long_text, result["optimized_prompt"])
            optimizer.print_comparison_result(comparison)
            
            # 保存比较结果
            combined_result = {
                "original_prompt": long_text,
                "optimization_result": result,
                "comparison_result": comparison
            }
            optimizer.save_optimization_result(combined_result)
        
        elif choice == "3":
            print("感谢使用，再见！")
            break
        else:
            print("无效的选择，请重新输入。")

if __name__ == "__main__":
    main()